package scfg;

import java.util.*;
import java.io.*;
import scfg.output.Display;
import scfg.output.RNAFormattedFile;
import scfg.utils.trainging_results.StatsCollector;

/**
 * The Class TrainingSet.
 */
public class TrainingSet {

	private Display output;

	/**
	 * The Enum GrammarType.
	 */
	public static enum GrammarType {

		/** The Knudsen hein. */
		KnudsenHein,
		/** The Pfold. */
		Pfold
	};

	/**
	 * Sets the grammar.
	 * 
	 * @param type
	 *            the type
	 * @param trainingSetFilename
	 *            the filename
	 */
	public Grammar setGrammar(GrammarType type, String trainingSetFilename) {
		switch (type) {
		case KnudsenHein:
			return new KHGrammar(trainingSetFilename, output);
		case Pfold:
			return new PfoldGrammar(trainingSetFilename, output);
		default:
		}
		return null;
	}

	/**
	 * Instantiates a new training set.
	 * 
	 * @param type
	 *            the type
	 */
	public TrainingSet(Display output) {
		this.output = output;
	}

	/**
	 * Reads a ".rna" file and pushes the contained sequence and structure onto
	 * the corresponding vectors. The file format should have the sequence on
	 * the first line and the structure on the second line. All other input will
	 * be ignored.
	 * 
	 * Multiple file format support is to come.
	 * 
	 * @param filename
	 *            - name of file to be read. Must be in above mentioned ".rna"
	 *            format.
	 * @param sequences
	 *            the sequences
	 * @param structures
	 *            the structures
	 */
	private void readRNAFile(String filename, List<RNAFormattedFile> trainingSet) {
		if (filename == null || trainingSet == null
				|| !filename.endsWith(".rna")) {
			output.eout("Illegal Parameters found in ReadRNAFile");
			return;
		}
		File f = new File(filename);
		if (f.exists()) {
			try {
				RNAFormattedFile rnaff = new RNAFormattedFile(f);
				if (rnaff.get("seq") != null && rnaff.get("nat") != null)
					trainingSet.add(rnaff);
			} catch (Exception e) {
				output.eout(e.getStackTrace().toString());
				e.printStackTrace();
			}
		} else
			output.eout("File does not exist: " + filename);
	}

	public int trainingSetFile(File f, GrammarType type, StatsCollector sc,
			boolean verbose) {
		Grammar grammar = setGrammar(type, f.getAbsolutePath());
		int count = 0;
		List<RNAFormattedFile> trainingSet = new ArrayList<RNAFormattedFile>();
		if (verbose)
			output.out("[Starting Training : " + f.getName() + "]");
		try {
			Scanner scan = new Scanner(f);
			while (scan.hasNextLine()) {
				readRNAFile(scan.nextLine(), trainingSet);
				count++;
			}
			scan.close();
			output.out("\tTotal of " + count + " files processed");
			if (count > 0) {
				grammar.train_grammar(grammar, trainingSet, verbose);
				if (sc != null)
					sc.add(grammar);
				// TODO: write to database
			} else
				output.out("Skipping: " + f.getName());
		} catch (Exception e) {
			output.eout(e.getStackTrace().toString());
			e.printStackTrace();
		}
		return count;
	}

	/*
	 * Can be used at anytime within the lifetime of the object to recalculate
	 * probabilities with a new training set. The provided directory is explored
	 * NON-recursively to analyse all files with extensions type "ext". The
	 * files are expected to have absolute paths to RNA files for training.
	 * 
	 * @param tdn - Training directory name. Recursively search from this
	 * directory.
	 * 
	 * @param ext - Extension of the training set files.
	 * 
	 * @param append - Appends results to current file if true;
	 */
	/**
	 * New training set.
	 * 
	 * @param dirName
	 *            the dir name
	 * @param ext
	 *            the ext
	 * @param append
	 *            the append
	 * @param verbose
	 *            the verbose
	 * @param graphs
	 *            the graphs
	 */
	public void trainingSetFileDirectory(String dirName, String ext, GrammarType type,
			boolean verbose) {
		StatsCollector sc = new StatsCollector();
		File base = new File(dirName);
		if (base.exists() && base.isDirectory()) {
			File[] contents = base.listFiles();
			for (File f : contents) {
				if (f.getName().endsWith(ext)) {
					trainingSetFile(f, type, sc, verbose);
				} else if (f.isDirectory()) {
					trainingSetFileDirectory(f.getAbsolutePath(), ext, type, verbose);
				}
			}
		} else {
			output.out("Directory not found: " + dirName);
		}
		if (verbose)
			output.out(sc.update());
		output.finalizeDisplay();
	}
	
	public void trainingSetDirectory(String dirName, GrammarType type,
			boolean verbose) {
		StatsCollector sc = new StatsCollector();
		String ext = ".rna";
		File base = new File(dirName);
		if (base.exists() && base.isDirectory()) {
			Grammar grammar = setGrammar(type, base.getAbsolutePath());
			List<RNAFormattedFile> trainingSet = new ArrayList<RNAFormattedFile>();
			File[] contents = base.listFiles();
			for (File f : contents) {
				if (f.getName().endsWith(ext)) {
					readRNAFile(f.getAbsolutePath(), trainingSet);
				} else if (f.isDirectory()) {
					trainingSetFileDirectory(f.getAbsolutePath(), ext, type, verbose);
				}
			}
			if (trainingSet.size() > 0) {
				grammar.train_grammar(grammar, trainingSet, verbose);
				if (sc != null)
					sc.add(grammar);
				// TODO: write to database
			} else
				output.out("Skipping: " + base.getName());
		} else {
			output.out("Directory not found: " + dirName);
		}
		if (verbose)
			output.out(sc.update());
		output.finalizeDisplay();
	}
	
	public static void main(String[] args){
		TrainingSet ts = new TrainingSet(new Display("TrainingSet"));
		File file = new File("/home/david/Documents/REU/RNA/training/5S/minList_20.txt");
		GrammarType type = GrammarType.Pfold;
		StatsCollector sc = null;
		boolean verbose = true;
		ts.trainingSetFile(file, type, sc, verbose);
		
		String dirName = "/home/david/Documents/REU/RNA/training/5S";
		String ext = ".txt";
		ts.trainingSetFileDirectory(dirName, ext, type, verbose);
	}
}
